#include "toroidal.h"
#include "cudafunc.h"
#include "math.h"

extern field Apsi;
extern field Bsca;

ptcs::ptcs(int N)
{
    this->npt = N;
    this->ppts = new particle[N];
    this->inside = new bool[N];
    this->pB = &Bsca;
    this->weights = new double[N];
    this->periods = new double[N];
    this->wdelta = new double[N];
}
ptcs::~ptcs()
{
    delete [] ppts;
}


double ptcs::histo(double rad)
{
    int count = 0;
    double r0 = 4.34;
    for (int i = 0; i < this->npt; i++)
    {
        double r = this->ppts[i].getcrd().r;
        double z = this->ppts[i].getcrd().z;
        if((r-r0)*(r-r0) + z*z < rad*rad)
            count++;
    }
    return (double)count / (double)this->npt;
}

void ptcs::output_realspace(const char* str)
{
    float* phiar = new float[this->npt];
    cudarandgen(phiar, this->npt,'u');
    FILE* fp = fopen(str,"w");
    fprintf(fp,"X,Y,Z\n");
    for(int i = 0; i < npt; i++)
    {
        vec cord = cyl2xyz(ppts[i].getcrd());
        fprintf(fp,"%f,%f,%f\n",cord.x, cord.y, cord.z);
    }
    fclose(fp);
}

void ptcs::output_poloidal(const char* str)
{
    FILE* fp = fopen(str,"w");
    fprintf(fp,"R,Z\n");
    for(int i = 0; i < npt; i++)
    {
        vec cord = cyl2xyz(ppts[i].getcrd());
        fprintf(fp,"%f,%f\n",sqrt(cord.x*cord.x+cord.y*cord.y), cord.z);
    }
    fclose(fp);
}

double ptcs::distribfunc(particle part)
{
    if (checkboundary(part.getcrd()) == false)
        return 0.0;
    double psip = Apsi.interp(part.getcrd());
    double pR = part.getcrd().r;
    return pR/Raxis*exp((psip/psiaxis)*3)/exp(3);
}

double ptcs::pol_dstrb(cylv crd)
{
    if (checkboundary_core(crd) == false)
        return 0.0;
    double psip = Apsi.interp(crd);
    return exp((psip/psiaxis)*3)/exp(3);
}

void ptcs::output_velocityspace(const char* str)
{
    float* phiar = new float[this->npt];
    cudarandgen(phiar, this->npt,'u');
    FILE* fp = fopen(str,"w");
    fprintf(fp,"vx,vy,vz\n");
    for(int i = 0; i < npt; i++)
    {
        if (inside[i] == false)
            continue;
        double Bfp = pB->interp(ppts[i].getcrd());
        double vptz = ppts[i].getvpara();
        double vperp2 = 2.0*ppts[i].pmu()*Bfp;
        double vptx = sqrt(vperp2)*cos(phiar[i]*2.0*pi);
        double vpty = sqrt(vperp2)*sin(phiar[i]*2.0*pi);
        fprintf(fp,"%f,%f,%f\n",vptx, vpty, vptz);
    }
    fclose(fp);
}

pfvec world::emfield(cylv crd)
{
    vec pE = this->E0->interp(crd);
    vec pB = this->B0->interp(crd);
    vec pBg = this->gB->interp(crd);
    vec pcb = this->cb->interp(crd);
    
    double cp = -crd.p;
    return {rotate(pE,cp), rotate(pB, cp),
    rotate(pBg, cp), rotate(pcb,cp)};
}

pfvec world::emfield_time(cylv crd, double tphase)
{
    int n = phies->n;
    double totalphase = (n*crd.p - tphase);
    vec pE = cos(totalphase)*this->E1i->interp(crd);
    pE = pE + sin(totalphase)*this->E1q->interp(crd);
    
    vec pB = this->B0->interp(crd);
    vec pBg = this->gB->interp(crd);
    vec pcb = this->cb->interp(crd);
    
    double cp = -crd.p;
    return {rotate(pE,cp), rotate(pB, cp),
    rotate(pBg, cp), rotate(pcb,cp)};
}

phsp world::push_rk2(pfvec fs, particle* pt, double dt, double vpara)
{
    double m = particle::m;
    double q = pt->q;
    double mu = pt->pmu();
    vec Bst = fs.B + q*vpara/m*fs.cb;
    double invBstn = sqrt(dot(fs.B,fs.B))/dot(Bst, fs.B);    
    vec dX = vpara*fs.B;
    vec curvdrift = q*vpara*vpara/m*fs.cb;
    vec graddrift = -pow(dot(fs.B,fs.B),-0.5)*cross(fs.B,fs.E + -1*mu/q*fs.gB);
    dX = invBstn *(dX + curvdrift + graddrift);
    double dvp =  dt*dot(Bst,fs.E+(-1*mu/q)*fs.gB)*q*invBstn/m;   
    return {dt*dX, dvp};
}

